package com.shujia.spark.sql

import org.apache.spark.sql.expressions.Window
import org.apache.spark.sql.{DataFrame, SparkSession}

object Demo6Window {
  def main(args: Array[String]): Unit = {
    /**
     * 窗口函数
     * row_number
     * max
     * min
     * sum
     * avg
     * count
     * rank
     * lag
     * lead
     */

    val spark: SparkSession = SparkSession
      .builder()
      .master("local")
      .appName("window")
      .config("spark.sql.shuffle.partitions", 1) //指定spark sql shuffle之后rdd的分区数
      .getOrCreate()
    import spark.implicits._
    import org.apache.spark.sql.functions._

    val studentDF: DataFrame = spark.read.json("data/students.json")

    /**
     * 1、max，min
     */
    //计算每个班级中每个人和班级年龄最大的学生相差几岁
    studentDF
      .withColumn("maxAge", max($"age") over Window.partitionBy($"clazz"))
      .withColumn("diff_age", $"maxAge" - $"age")
      .show()

    /**
     * 2、sum
     *
     * expr: 将一个sql表达式转换成列对象
     */
    val scoreDF: DataFrame = spark.read
      .format("csv")
      .option("sep", ",")
      .schema("id STRING,cid STRING ,sco DOUBLE")
      .load("data/score.txt")

    scoreDF.show()

    //组内求和
    //计算每个学生每个科目的分数占该学生总分的比例
    scoreDF
      .withColumn("sumSco", sum($"sco") over Window.partitionBy($"id"))
      .withColumn("p", concat(round($"sco" / $"sumSco" * 100, 1), expr("'%'")))
      .show()

    //累积求和
    scoreDF
      .withColumn("sum", sum($"sco") over Window.partitionBy($"id").orderBy($"sco"))
      .show()

    /**
     * 3、count
     */
    //3、统计每科都及格的学生
    val subjectDF: DataFrame = spark
      .read
      .format("csv")
      .option("sep", ",")
      .schema("cid STRING ,cname STRING , sumSco DOUBLE")
      .load("data/subject.txt")

    scoreDF
      .join(subjectDF, "cid") //关联科目表
      .where($"sco" >= $"sumSco" * 0.6) //取出及格和的分数
      .withColumn("num", count($"id") over Window.partitionBy($"id")) //统计及格的数量
      .where($"num" === 6)
      .show()

    /**
     * 4、lag,lead
     */
    //统计每个班级中每个学生的总分和前面一个名次相差的分数
    scoreDF
      .groupBy($"id")
      .agg(sum($"sco") as "sumSco") //计算总分
      .join(studentDF, "id") //关联学生表获取班级
      .withColumn("next_sum_sco", lead($"sumSco", 1) over Window.partitionBy($"clazz").orderBy($"sumSco"))
      .withColumn("diff_sco", $"next_sum_sco" - $"sumSco")
      .show()

  }

}
